#include <iostream>
#include <fstream>
#include <sstream>
#include "math.h"

#include "cv.h"
#include "highgui.h"
#include "common_defs.h"

#include "PointsDatabase.h"
#include "ImageDepthList.h"

#include "PointsOfGradUI.h"
#include "PointsUI.h"
#include "FigOrCatDef.h"
#include "GradOrCatDef.h"
#include "FigAreaCatDef.h"
#include "DimCatDef.h"
#include "SpatialFreqCatDef.h"
#include "NumSegCatDef.h"
#include "ComplexityCatDef.h"


PointsOfGradUI::PointsOfGradUI( ImageDepthList *imlist, PointsDatabase *pts, float tolDeg ) 
	: m_imUI_temp(cvSize(640,640),cvSize(300,300),cvSize(200,200) )
{
	m_imlist = imlist;
	m_pts = pts;
	
	m_tolDeg = tolDeg;
	m_tol = tolDeg * (float)PI / 180.f; // 5 degrees

	m_localPxRadius = 25.f;

	m_numGrads = cvCeil( 180.f / (2*tolDeg) );
	m_aPtsGrad = new list< PointOfGrad* >[m_numGrads];
	
	m_ptsUI = new PointsUI();
	m_imUI = new PointsOfImageUI();

	cvInitFont( &m_font, CV_FONT_HERSHEY_DUPLEX, .5f, .5f, 0, 1, CV_AA );

	m_checkBoxHeight = 30;

	m_imUI_params.bRedrawNeeded = false;
}

PointsOfGradUI::~PointsOfGradUI()
{
	delete m_ptsUI;
	delete m_imUI;
}

void deletePtsOfGrad( list< PointOfGrad* > &ptsOfGrad )
{
	while( !ptsOfGrad.empty() )
	{
		delete ptsOfGrad.back();
		ptsOfGrad.pop_back();
	}
}

void PointsOfGradUI::categorizePtsOfGrad( CatMap &catPts, list< PointOfGrad* > &ptsOfGrad, const MouseCallbackParams &p )
{
	while( !catPts.empty() ) {
		delete catPts.begin()->second;
		catPts.erase( catPts.begin() );
	}

	for( list<PointOfGrad*>::iterator pgiter = ptsOfGrad.begin(); pgiter != ptsOfGrad.end(); pgiter++ )
	{
		string catvals;
		string val;
		// iterate through the categories
		for( CatUIList::const_iterator cuit = p.cats.begin(); cuit != p.cats.end(); cuit++ )
		{
			if( !cuit->enabled ) continue;

			ImageDepthList::iterator imdIter = m_imlist->find( (*pgiter)->imIter->first );
			ImageGroupData imdata;

			cuit->def->categorize( imdata, *imdIter->second, *(*pgiter)->iter, val, m_localPxRadius );

			if( !catvals.empty() ) catvals += ",";
			catvals += cuit->def->key() + "=" + val;
		}
		if( catPts[catvals]==NULL ) catPts[catvals] = new CategorizedPts( catvals );
		catPts[ catvals ]->pts.push_back( *pgiter );
		catPts[ catvals ]->imNames.insert( (*pgiter)->imIter->first );
	}

	int lineOffset = 0;
	for( CatMap::iterator cmit = catPts.begin(); cmit != catPts.end(); cmit++ )
	{
		lineOffset += 1; // make room for text describing category
		
		if( catPts.size() > 1 ) cmit->second->lineOffset = lineOffset;
		else  cmit->second->lineOffset = 0;

		lineOffset += cvCeil( (double) cmit->second->pts.size() / p.winSizePnl.width );
	}
}
void PointsOfGradUI::drawCheckBoxes( IplImage* dispIm, MouseCallbackParams &p )
{
	int h = m_checkBoxHeight;
	CvRect roi = cvRect( 0, dispIm->height - h, dispIm->width, h );
	drawRect( dispIm, roi, cvScalarAll(180), -1 );
	cvLine( dispIm, cvPoint(0,roi.y), cvPoint(roi.width,roi.y), cvScalarAll(255), 1, CV_AA, 0 );
	
	int cursor = 15;
	CvRect chkbox = cvRect( cursor, dispIm->height - 3*h/4, h/2,h/2);
	for( CatUIList::iterator cuit = p.cats.begin(); cuit != p.cats.end(); cuit++ )
	{
		chkbox.x = cursor;
		cuit->chkbox = chkbox;

		drawRect( dispIm, chkbox, cvScalarAll(140), -1 );
		cvLine( dispIm, cvPoint(chkbox.x-1,chkbox.y-1), cvPoint(chkbox.x-1,chkbox.y+chkbox.height), cvScalarAll(50) );
		cvLine( dispIm, cvPoint(chkbox.x-1,chkbox.y-1), cvPoint(chkbox.x+chkbox.width,chkbox.y-1), cvScalarAll(50) );
		cursor += h;

		if( cuit->enabled )
		{
			cvLine( dispIm, cvPoint(chkbox.x-1,chkbox.y-1), cvPoint(chkbox.x+chkbox.width,chkbox.y+chkbox.height), cvScalarAll(0), 2 );
			cvLine( dispIm, cvPoint(chkbox.x-1,chkbox.y+chkbox.height), cvPoint(chkbox.x+chkbox.width,chkbox.y-1), cvScalarAll(0), 2 );	
		}

		cvPutText( dispIm, cuit->def->key().c_str(), cvPoint( cursor, dispIm->height - h/2 + 4 ), &m_font, CV_RGB(0,0,55));
		cursor += cuit->def->key().length() * 10 + 20;
	}
}
void PointsOfGradUI::drawCategorizedPtsOfGrad( IplImage* dispIm, CatMap &catPts, MouseCallbackParams &p )
{
	for( CatMap::iterator cmit = catPts.begin(); cmit != catPts.end(); cmit++ )
	{
		CvRect headingRoi = 
			cvRect( 0, (cmit->second->lineOffset - p.lineOffset-1)*p.panelSizePx.height + p.panelSizePx.height/8, 
			p.winSizePx().width, p.panelSizePx.height * 6 / 8 );

		if( catPts.size() > 1 )
		{
			drawRect( dispIm, headingRoi, CV_RGB(250,250,255), -1 );
			CvPoint txtPt = cvPoint( 50, headingRoi.y + p.panelSizePx.height / 2 );
			{
				stringstream strOut;
				strOut << cmit->first.c_str() << " : " << cmit->second->pts.size() << " matches from ";
				strOut << cmit->second->imNames.size() << " images";
				
				cvPutText( dispIm, strOut.str().c_str(), txtPt , &m_font, CV_RGB(0,0,255) );
				
				printf("\t%d", cmit->second->pts.size() );
			}
		}

		int i = 0;
		for( list< PointOfGrad* >::iterator pgiter = cmit->second->pts.begin(); pgiter != cmit->second->pts.end(); pgiter++, i++ )
		{
			int lineoffset = cmit->second->lineOffset - p.lineOffset;
			CvPoint pt = 
				cvPoint( 
					(i%p.winSizePnl.width)*p.panelSizePx.width + p.panelSizePx.width/2, 
					(i/p.winSizePnl.width + lineoffset )*p.panelSizePx.height + p.panelSizePx.height/2 );

			const CvRect dispRoi0 = cvRect( ( i%p.winSizePnl.width)*p.panelSizePx.width, 
				( i/p.winSizePnl.width + lineoffset)*p.panelSizePx.height, p.panelSizePx.width, p.panelSizePx.height );
			CvRect dispRoi = dispRoi0;

			if( ( i/p.winSizePnl.width + lineoffset) < 0 || ( i/p.winSizePnl.width + lineoffset) >= p.winSizePnl.height )
			{
				continue;
			}
			dispRoi.width = (*pgiter)->smImage->width;
			dispRoi.height = (*pgiter)->smImage->height;

			if( (*pgiter)->imRoi.x == 0 )
			{
				dispRoi.x += p.panelSizePx.width - dispRoi.width;
			}
			if( (*pgiter)->imRoi.y == 0 )
			{
				dispRoi.y += p.winSizePnl.height - dispRoi.height;
			}
			cvSetImageROI( dispIm, dispRoi );
			
			cvCopy( (*pgiter)->smImage, dispIm );

			cvResetImageROI( dispIm );

			float grad = (float)(*pgiter)->iter->gradRad;
			drawLineAtAngle( dispIm, pt, grad-PI/2, 5.f, p.panelSizePx.width/2-10.f, CV_RGB(255,0,0), 1 );
			drawLineAtAngle( dispIm, pt, grad+PI/2, 5.f, p.panelSizePx.width/2-10.f, CV_RGB(255,0,0), 1 );
			drawLineAtAngle( dispIm, pt, grad-PI/2, p.panelSizePx.width/2-5.f, p.panelSizePx.width/2, CV_RGB(0,0,0), 1 );
			drawLineAtAngle( dispIm, pt, grad+PI/2, p.panelSizePx.width/2-5.f, p.panelSizePx.width/2, CV_RGB(0,0,0), 1 );
			drawLineAtAngle( dispIm, pt, grad-PI/2, p.panelSizePx.width/2-10.f, p.panelSizePx.width/2-5.f, cvScalarAll(255), 1 );
			drawLineAtAngle( dispIm, pt, grad+PI/2, p.panelSizePx.width/2-10.f, p.panelSizePx.width/2-5.f, cvScalarAll(255), 1 );

			if( (*pgiter)->selected )
			{
				cvRectangle( dispIm, 
					cvPoint( dispRoi0.x, dispRoi0.y ),
					cvPoint( dispRoi0.x + dispRoi0.width - 2, dispRoi0.y + dispRoi0.height - 2 ),
					CV_RGB(255,255,0), 2, 4 );
			}
		}
	}
	printf("\n");

	drawCheckBoxes( dispIm, p );
}
void PointsOfGradUI::deleteSelectedPtsOfGrad( list< PointOfGrad* > &ptsOfGrad, ImageDepthList &imlist, CvSize winLayout, CvSize winSize, IplImage* dispIm, int lineOffset )
{
	for( list< PointOfGrad*>::iterator pgiter = ptsOfGrad.begin(); pgiter != ptsOfGrad.end();  ) {
		if( (*pgiter)->selected )
		{
			(*pgiter)->imIter->second.erase( (*pgiter)->iter );
			pgiter = ptsOfGrad.erase( pgiter );
		}
		else
		{
			pgiter++;
		}
	
	}
}
void unselectPtsOfGrad( list< PointOfGrad* > &ptsOfGrad, ImageDepthList &imlist, CvSize winLayout, CvSize winSize, IplImage* dispIm, int lineOffset )
{
	for( list< PointOfGrad*>::iterator pgiter = ptsOfGrad.begin(); pgiter != ptsOfGrad.end(); pgiter++ ) {
		(*pgiter)->selected = false;
	}
}
inline unsigned char mirrorBits(unsigned char value)
{
    unsigned char ret = 0;
    for( int i = 0; i < 8; i++ )
	{
        if ((value & (unsigned char)(1 << i)) != 0)
		{
			ret += (unsigned char)(1 << (7 - i));
		}
	}
    return ret;
}
void ConvertSeg( const void* _src, void* _dst )
{
	const unsigned char* src = (const unsigned char*) _src;
	unsigned char* dst = (unsigned char*) _dst;
	dst[0] = mirrorBits(src[0]);
	dst[1] = dst[0];
	dst[2] = dst[0];
}

void PointsOfGradUI::run()
{
	CatMap catmap;
	MouseCallbackParams p( this, cvSize(20,16), cvSize(50,50), &catmap );

	p.cats.push_back( CategoryUIInfo( new FigOrCatDef() ) );
	p.cats.push_back( CategoryUIInfo( new GradOrCatDef() ) );
	p.cats.push_back( CategoryUIInfo( new FigAreaCatDef() ) );
	p.cats.push_back( CategoryUIInfo( new SpatialFreqCatDef() ) );
	//p.cats.push_back( CategoryUIInfo( new DimCatDef() ) );
	p.cats.push_back( CategoryUIInfo( new NumSegCatDef() ) );

	p.cats.push_back( CategoryUIInfo( new ComplexityCatDef() ) );

	CvScalar bg = cvScalarAll(120);
	IplImage* dispIm = cvCreateImage( p.winSizePx(), IPL_DEPTH_8U, 3 );
	cvSet( dispIm, bg );
	
	float desiredGradDeg = 0.f;
	if(0)
	{
		cout << "Receptive field orientation (in degrees): ";
		cin >> desiredGradDeg;
		cin.ignore();
	}
	
	cvNamedWindow( "main", 1);
	cvShowImage("main", dispIm );
	cvSetMouseCallback( "main", PointsOfGradUI::mouseCallback, &p );
	cvWaitKey( 10 );

	m_desiredGrad = ( desiredGradDeg + 90.f ) * PI / 180.f; 
	refreshPoints( p );
	//categorizePtsOfGrad( catmap, m_aPtsGrad[p.gradIndex], p );

	//list< PointOfGrad* > ptsOfGrad; // points of the desired gradient

	int key = 0;
	do
	{
		//printf("Finding points with %.2f\xF1%.2f\xA7 orientation...\n", desiredGradDeg, m_tolDeg );
		printf("%.2f\xA7\t", desiredGradDeg );
		
		cvSetImageROI( dispIm, cvRect(0,0,p.panelSizePx.width, p.panelSizePx.width));
		cvZero(dispIm);
		drawLineAtAngle( dispIm, cvPoint( p.panelSizePx.width/2, p.panelSizePx.height/2), m_desiredGrad-PI/2, 0, p.panelSizePx.width+p.panelSizePx.height, cvScalarAll(200), 1 );
		drawLineAtAngle( dispIm, cvPoint( p.panelSizePx.width/2, p.panelSizePx.height/2), m_desiredGrad+PI/2, 0, p.panelSizePx.width+p.panelSizePx.height, cvScalarAll(200), 1 );
		cvResetImageROI( dispIm );

		//cout << "There are " << m_aPtsGrad[p.gradIndex].size() << " matching points." << endl;

		key = -1;
		while( key == -1 )
		{
			if( p.catUpdated || p.redrawDisplayImage )
			{
				cvScale( dispIm, dispIm, .5, 50);
				cvNamedWindow( "main", 1);
				cvShowImage("main", dispIm );
				key = cvWaitKey(50);
			}
			if( p.catUpdated )
			{
				categorizePtsOfGrad( catmap, m_aPtsGrad[p.gradIndex], p );
				p.catUpdated = false;
				p.redrawDisplayImage = true;
				
				for( CatMap::iterator cmit = p.catPts->begin(); cmit != p.catPts->end(); cmit++ )
				{
					printf( "\t%s", cmit->first.c_str() );
				}
				printf("\n");
			}
			if( p.redrawDisplayImage )
			{
				cvSet( dispIm, bg );
				drawCategorizedPtsOfGrad( dispIm, catmap, p );

				if( p.bRedrawLastClicked && p.lastClicked != NULL )
				{
					ImageDepthList::iterator imdIter = m_imlist->find( p.lastClicked->imIter->first );
					ImageGroupData imdata( imdIter->second );

					cvNamedWindow( "seg", 1);
					IplImage* seg = cvCreateImage( cvGetSize(imdata.getImage()), 8, 3);
					cvCvtColor( imdata.getSegm(), seg, CV_GRAY2BGR );

					applyToImage( seg, seg, &ConvertSeg );
					cvConvertScale( seg, seg, -1, 255);

					float frad = (float)p.lastClicked->iter->figRad ;
					float grad = (float)p.lastClicked->iter->gradRad ;
					CvPoint pt = p.lastClicked->iter->loc;
					
					IplImage* im[2] = { seg, imdata.image };
					for( int i = 0; i < 2; i++ )
					{
						drawLineAtAngle( im[i], pt, frad-PI/2, 5.f, 25.f, CV_RGB(255,255,0), 1 );
						drawLineAtAngle( im[i], pt, frad+PI/2, 5.f, 25.f, CV_RGB(255,255,0), 1 );
						drawLineAtAngle( im[i], pt, grad-PI/2, 7.f, 20.f, CV_RGB(255,0,0), 1 );
						drawLineAtAngle( im[i], pt, grad+PI/2, 7.f, 20.f, CV_RGB(255,0,0), 1 );
						cvCircle( im[i], pt, 26, CV_RGB(255,255,0),1,CV_AA);
					}

					cvShowImage("seg", seg );

					cvNamedWindow( "im", 1);
					cvShowImage("im", imdata.getImage() );

					cvReleaseImage( &seg );
				}
				p.redrawDisplayImage = false;
			}
			if( m_imUI_params.bRedrawNeeded )
			{
				m_imUI->drawWindow( m_imUI_params, m_imUI_temp );
				display( "Image View", m_imUI_temp.dispImg, 1 );
			}

			list< PointOfGrad* >::iterator * modPtIter = p.modPtIter;
			p.modPtIter = NULL;
			if( modPtIter )
			{
				//m_ptsUI->run(
			}

			cvNamedWindow( "main", 1);
			cvShowImage("main", dispIm );
			key = cvWaitKey(50);

			switch( key )
			{
			case 2555904: // right arrow
				p.gradIndex--;
				if( p.gradIndex < 0 ) p.gradIndex = m_numGrads - 1;
				desiredGradDeg -= 2*m_tolDeg;
				
				categorizePtsOfGrad( catmap, m_aPtsGrad[p.gradIndex], p );
				p.lineOffset = 0;
				p.redrawDisplayImage = true;
				break;
			case 2424832: // left arrow
				p.gradIndex = (p.gradIndex+1)%m_numGrads;
				desiredGradDeg += 2*m_tolDeg;
				
				categorizePtsOfGrad( catmap, m_aPtsGrad[p.gradIndex], p );
				p.lineOffset = 0;
				p.redrawDisplayImage = true;
				break;
			case 2621440: // down arrow
				//if( ( m_aPtsGrad[p.gradIndex].size()/p.winSizePnl.width - p.lineOffset) >= p.winSizePnl.height )
				//if( !catmap.empty() && 
				//	( m_aPtsGrad[p.gradIndex].size()/p.winSizePnl.width + catmap.rend()->second->lineOffset  - p.lineOffset ) >= p.winSizePnl.height )
				{
					p.lineOffset += p.winSizePnl.height/2;
					p.redrawDisplayImage = true;
				}
				break;
			case 2490368: // up arrow
				p.lineOffset -= p.winSizePnl.height/2 ;
				if( p.lineOffset < 0 ) p.lineOffset = 0;
				p.redrawDisplayImage = true;
				break;
			case 43: // plus sign
				p.winSizePnl.width += 5;
				cvReleaseImage(&dispIm);
				dispIm = cvCreateImage( p.winSizePx(), IPL_DEPTH_8U, 3 );
				p.redrawDisplayImage = true;
				break;
			case 3014656: // delete key
				deleteSelectedPtsOfGrad( m_aPtsGrad[p.gradIndex], *m_imlist, p.winSizePnl, p.panelSizePx, dispIm, p.lineOffset );
				
				m_pts->save();
				refreshPoints( p );
				p.catUpdated = true;
				break;
			case 'r': case 'R': case 7602176: // F5 - refresh
				refreshPoints( p );
				p.catUpdated = true;
				break;
			case 'f': case 'F':
				categorizePtsOfGrad( catmap, m_aPtsGrad[p.gradIndex], p );
				p.redrawDisplayImage = true;
				break;
			case -1:
				break;
			default:
				printf("key = %d\n", key );
				break;
			}
		}
	} while( key != 27 );

	cvReleaseImage( &dispIm );

	for( int gi = 0; gi < m_numGrads; gi++ )
	{
		deletePtsOfGrad( m_aPtsGrad[gi] );
	}
	while( !catmap.empty() ) {
		delete catmap.begin()->second;
		catmap.erase( catmap.begin() );
	}
	
	while( !p.cats.empty() ) {
		p.cats.back().deleteThis();
		p.cats.pop_back();
	}
}

void PointsOfGradUI::refreshPoints( MouseCallbackParams& p )
{
	for( size_t gi = 0; gi < m_numGrads; gi++ )
	{
		// take care of memory
		for( list< PointOfGrad* >::iterator iter = m_aPtsGrad[gi].begin(); 
			  iter != m_aPtsGrad[gi].end(); iter++ )
		{
			delete *iter;
		}
		m_aPtsGrad[gi].clear();
	}
	for( PointsDatabase::iterator imIter = m_pts->begin(); imIter != m_pts->end(); imIter++ )
	{
		ImageDepthList::iterator imdIter = m_imlist->find( imIter->first );
		if( imdIter == m_imlist->end() )
		{
			printf("Unable to lookup %s!\n", imIter->first.c_str() );
			assert( imdIter != m_imlist->end() );
			continue;
		}
		ImageGroupData imdata;
		imdata.image = cvLoadImage( imdIter->second->pathImage.c_str() );
		if( !imdata.image )
		{
			cerr << "Unable to open file " << imIter->first << "!" << endl;
			continue;
		}

		for( PointsDatabase::PointsInImage::iterator piter = imIter->second.begin();
			piter != imIter->second.end(); piter++ )
		{
			float diff = fmodulus<float>( piter->gradRad + m_tol - m_desiredGrad, PI );
			size_t gi = diff / (2*m_tol); //radianDifference<float>( , PI ) + m_tol;

			m_aPtsGrad[gi].push_back(  new PointOfGrad( imIter, piter, imdata, p.panelSizePx ) );
		}
		
		imdata.deleteAll();
	}
}


void PointsOfGradUI::mouseCallback( int evt, int x, int y, int flags, void* vparam )
{
	if( !evt ) return;
	MouseCallbackParams* p = (MouseCallbackParams*)vparam;
	PointsOfGradUI* pogui = p->pogui;

	if( y <= p->winSizePx().height - pogui->m_checkBoxHeight )
	{
		int px = x / p->panelSizePx.width;
		int py = y / p->panelSizePx.height + p->lineOffset;
		
		CatMap::iterator inCat = p->catPts->end();
		for( CatMap::iterator cmit = p->catPts->begin(); cmit != p->catPts->end(); cmit++ )
		{
			if( cmit->second->lineOffset <= py ) {
				inCat = cmit;
			}
		}
		if( inCat == p->catPts->end() ) return;

		py -= inCat->second->lineOffset;

		int num = py * p->winSizePnl.width + px;

		if( num >= inCat->second->pts.size() ) return;

		list< PointOfGrad* >::iterator iter = inCat->second->pts.begin();
		for( int i = 0; i < num && iter != inCat->second->pts.end(); i++, iter++ )
		{}
		if( iter == inCat->second->pts.end() ) return;

		switch( evt )
		{
		case CV_EVENT_LBUTTONDOWN:
			{
				(*iter)->selected = !(*iter)->selected;
				if( p->lastClicked != *iter ) p->bRedrawLastClicked = true;
				p->lastClicked = *iter;
				p->redrawDisplayImage = true;
			}
			break;
		case CV_EVENT_LBUTTONDBLCLK:
			{
				p->modPtIter = &iter;
			}
			break;
		case CV_EVENT_RBUTTONDOWN:
			{
				display( "Image View", pogui->m_imUI_temp.dispImg, 1 );
				cvSetMouseCallback( "Image View", PointsOfImageUI::mouseCallback, &pogui->m_imUI_params );

				ImageDepthList::const_iterator imdIter = pogui->m_imlist->find( (*iter)->imIter->first );
				pogui->m_imUI_params.pts = pogui->m_pts;
				pogui->m_imUI_params.chosenPt = &(*(*iter)->iter);

				pogui->m_imUI->prepare( imdIter, pogui->m_imUI_params, pogui->m_imUI_temp );
				pogui->m_imUI->drawWindow( pogui->m_imUI_params, pogui->m_imUI_temp );
				display( "Image View", pogui->m_imUI_temp.dispImg, 1 );

				//(*iter)->selected = !(*iter)->selected;
				//pogui->m_imUI;
			}
			break;
		}
	}
	else
	{
		if( evt == CV_EVENT_LBUTTONDOWN )
		{
			for( CatUIList::iterator cuit = p->cats.begin(); cuit != p->cats.end(); cuit++ )
			{
				if( x > cuit->chkbox.x && x <= cuit->chkbox.x + cuit->chkbox.width )
				{
					cuit->enabled = !cuit->enabled;
					p->catUpdated = true;
				}
			}
		}
	}
}